# -*- coding: utf-8 -*-
# @Time    : 2021/3/24 9:24
# @Author  : huangwei
# @File    : server.py
# @Software: PyCharm
import os
import cv2
import json
import base64
import numpy as np
from PIL import Image
from flask import request, Flask
from wsgiref.simple_server import make_server

from config import args
from table import get_cell
from method import create_dir, delete_dir, getByte
from word import TextDetector, TextRecognizer
from signature import sort_box_to_group, cut_name

### 加载模型
text_detector = TextDetector(args)
text_recognizer = TextRecognizer(args)
WEIGHT_PATH = "models/table.weights"
tableNet = cv2.dnn.readNetFromDarknet(WEIGHT_PATH.replace('.weights', '.cfg'), WEIGHT_PATH)

app = Flask(__name__)


@app.route("/", methods=['POST'])
def get_frame():
    upload_file = request.get_data()
    req = json.loads(upload_file)

    if upload_file:
        filepath = req['filepath']
        print("filepath is:", filepath)
        name = os.path.basename(filepath)

        img_str = req['image']  # 得到unicode的字符串
        # img_decode_ = img_str.encode('ascii')  # 从unicode变成ascii编码
        img_decode = base64.b64decode(img_str)  # 解base64编码，得图片的二进制
        img_np_ = np.frombuffer(img_decode, np.uint8)
        img = cv2.imdecode(img_np_, cv2.COLOR_RGB2BGR)  # 转为opencv格式

        tmp_path = "./temp_path"
        create_dir(tmp_path)

        img_path = "%s/%s" % (tmp_path, name)
        height, width = img.shape[0:2]
        if height > 2000 or width > 2000:
            if height > width:
                new_height = 2000
                new_width = int(new_height / height * width)
            else:
                new_width = 2000
                new_height = int(new_width / width * height)
            img = cv2.resize(img, (new_width, new_height))

        cv2.imwrite(img_path, img)
        cv2.imencode('.png', img)[1].tofile(img_path)

        cols = req['cols']
        print(cols)

        img_end = {'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif'}
        if os.path.isfile(img_path) and os.path.splitext(img_path)[-1][1:].lower() in img_end:
            # 检测出图片中文字出现的位置
            img = cv2.imread(img_path)
            dt_boxes = text_detector(img)

            # 获取表格信息，返回每一条横线和每一个空格
            img2 = Image.open(img_path).convert('RGB')
            boxes, row_lines = get_cell(img2, tableNet, prob=0.5, alph=10, col=30)

            from test import draw_box
            for box in boxes:
                draw_box(img, box)
            cv2.imwrite("pic/box.png", img)
            print("box numbers:", len(boxes))

            # 使用检测出的表格线划分每一行检测出的表格和文字格
            group_boxes, group_det_boxes = sort_box_to_group(boxes, dt_boxes, row_lines)

            output_path = "./output"
            create_dir(output_path)

            # 对每一行的box进行识别和剪切
            for i in range(len(group_boxes)):
                cut_name(img_path, group_boxes[i], group_det_boxes[i], output_path, i, text_recognizer, cols)
            print("图片裁剪成功！")

            response_list = []
            # 将output文件夹中的图片返回，再清空temp_path,output文件夹
            for file in os.listdir(output_path):
                path = os.path.join(output_path, file)
                file_str = getByte(path)
                response = {'name': file, 'image': file_str}
                response_list.append(response)

            res = json.dumps(response_list)

            # delete_dir(output_path)
            # delete_dir(tmp_path)
            return res

        else:
            return "this is not a picture"
    else:
        return 'failed to get signature!'


if __name__ == "__main__":
    # app.run("127.0.0.1", port=6060)
    try:
        server = make_server("", port=6060, app=app)
        print("Running on http://localhost:6060/")
        server.serve_forever()
    except:
        pass
